import os
import json
import requests
from datetime import datetime
import time
import zipfile
from .constants import DCE_BASE_URL, DCE_DOWNLOAD_URLS, DCE_DATA_PATH, DCE_FUTURES

class DCEDownloader:
    def __init__(self):
        self.base_url = DCE_BASE_URL
        self.session = requests.Session()
        self.headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
        }
        self.session.headers.update(self.headers)
        
    def _create_year_folder(self, year):
        """创建年份文件夹"""
        year_path = os.path.join(DCE_DATA_PATH, str(year))
        if not os.path.exists(year_path):
            os.makedirs(year_path)
        return year_path
        
    def _extract_zip(self, zip_path, extract_path):
        """解压zip文件，处理中文编码"""
        try:
            with zipfile.ZipFile(zip_path, 'r') as zip_ref:
                # 处理zip文件中的中文编码
                for file in zip_ref.namelist():
                    try:
                        # 尝试不同的编码
                        for encoding in ['gbk', 'gb2312', 'gb18030', 'utf-8']:
                            try:
                                file_name = file.encode('cp437').decode(encoding)
                                break
                            except:
                                continue
                        
                        # 解压单个文件
                        data = zip_ref.read(file)
                        out_path = os.path.join(extract_path, file_name)
                        
                        # 确保输出路径存在
                        os.makedirs(os.path.dirname(out_path), exist_ok=True)
                        
                        with open(out_path, 'wb') as f:
                            f.write(data)
                    except Exception as e:
                        print(f"解压文件 {file} 失败: {str(e)}")
                        
            os.remove(zip_path)  # 解压后删除zip文件
            return True
        except Exception as e:
            print(f"解压失败: {str(e)}")
            return False
            
    def _download_file(self, url, save_path):
        """下载单个文件"""
        try:
            full_url = self.base_url + url if not url.startswith('http') else url
            response = self.session.get(full_url)
            
            if response.status_code == 200:
                # 确保文件夹存在
                os.makedirs(os.path.dirname(save_path), exist_ok=True)
                
                with open(save_path, 'wb') as f:
                    f.write(response.content)
                    
                # 如果是zip文件，自动解压
                if save_path.endswith('.zip'):
                    extract_path = os.path.dirname(save_path)
                    if self._extract_zip(save_path, extract_path):
                        print("文件已解压")
                    else:
                        return False
                return True
            else:
                print(f"下载失败: {response.status_code}")
                return False
                
        except Exception as e:
            print(f"下载出错: {str(e)}")
            return False
            
    def download_history(self, start_year=2010, end_year=None):
        """下载历史数据"""
        if end_year is None:
            end_year = datetime.now().year
            
        for year in range(start_year, end_year + 1):
            year_str = str(year)
            if year_str not in DCE_DOWNLOAD_URLS:
                print(f"{year}年数据链接不存在")
                continue
                
            print(f"\n开始下载 {year} 年数据...")
            year_path = self._create_year_folder(year)
            
            # 下载该年份的所有品种数据
            for code, name in DCE_FUTURES.items():
                if code not in DCE_DOWNLOAD_URLS[year_str]:
                    continue
                    
                download_url = DCE_DOWNLOAD_URLS[year_str][code]
                ext = os.path.splitext(download_url)[1]
                
                # 简化文件名，避免中文导致的编码问题
                file_name = f"{code}{ext}"
                save_path = os.path.join(year_path, file_name)
                
                print(f"正在下载 {code} {year}年数据...")
                if self._download_file(download_url, save_path):
                    print(f"已保存: {save_path}")
                
                time.sleep(2)
        
        print("\n所有数据下载完成！")

def download_dce_history(start_year=2010, end_year=None):
    """下载大连商品交易所历史数据的便捷函数"""
    try:
        downloader = DCEDownloader()
        downloader.download_history(start_year, end_year)
    except Exception as e:
        print(f"\n运行出错: {str(e)}")
        raise

if __name__ == "__main__":
    # 下载2010年至今的所有数据
    download_dce_history()
